from spinup import ppo_pytorch as ppo
from spinup import ddpg_pytorch as ddpg
from spinup.algos.pytorch.td3.td3 import td3 as td3_pytorch
from spinup.utils.test_policy import load_policy_and_env, run_policy
import torch
import gym

TRAIN = 1
test = 'test4'
env = lambda : gym.make(f'ur5controllerdualPPO{test}-v2')

if TRAIN:

    if test == 'test1':
        ### tset1 targetpos=[-0.35 0.6 0.Dual-UR5-IK-Traj-Plan] size=[0.Dual-UR3-Traj-Plan .05 .05]
        ac_kwargs = dict(hidden_sizes=[128,128,128], activation=torch.nn.ReLU)
        logger_kwargs = dict(output_dir='test1', exp_name='Dual_UR3')

    elif test == 'test2':
        ### tset2 targetpos=[-0.35 0.6 0.7] size=[0.Dual-UR3-Traj-Plan .05 .05]
        ac_kwargs = dict(hidden_sizes=[128,128,128], activation=torch.nn.ReLU)
        logger_kwargs = dict(output_dir='test2', exp_name='Dual_UR3')

    elif test == 'test3':
        ### tset3 targetpos=[-0.35 0.6 0.5] size=[0.One-Arm-UR3-Test .05 .05]
        ac_kwargs = dict(hidden_sizes=[128,128,128], activation=torch.nn.ReLU)
        logger_kwargs = dict(output_dir='test3', exp_name='Dual_UR3')

    elif test == 'test4':
        ### tset4 targetpos=[-0.35 -0.6 0.6] size=[0.One-Arm-UR3-Test .05 .05]
        ac_kwargs = dict(hidden_sizes=[128,128,128], activation=torch.nn.ReLU)
        logger_kwargs = dict(output_dir='test4', exp_name='Dual_UR3')

    ppo(env, ac_kwargs=ac_kwargs, logger_kwargs=logger_kwargs, seed=0,
        steps_per_epoch=5000, epochs=300, gamma=0.99, clip_ratio=0.2, pi_lr=3e-4,
        vf_lr=1e-3, train_pi_iters=80, train_v_iters=80, lam=0.97, max_ep_len=500,
        target_kl=0.05, save_freq=10)

else:
    _, get_action = load_policy_and_env(test)
    env_test = gym.make(f'ur5controllerdualPPO{test}-v2')
    run_policy(env_test, get_action, test)
